"""3v3坦克对抗环境实现"""

import math
from pathlib import Path

import addict
import toml
from matplotlib import pyplot as plt
from matplotlib.patches import Circle

from ._typing import Action, Info, Observation, Reward
from .soldier import Soldier
from .tank import Tank
from .team import Team
from .utils import atan, compute_distance

__all__ = ["ToyEnv", "config"]

plt.rcParams.update({"font.sans-serif": "Microsoft YaHei"})
plt.switch_backend("TkAgg")

config = addict.Dict(toml.load(Path(__file__).parent / "config.toml"))


class ToyEnv:
    """简易游戏环境。"""

    window: plt.Axes

    def __init__(self, render_mode: str | None = None, render_fps: int = 5) -> None:
        self.config = config
        self.render_mode = render_mode
        self.render_fps = render_fps
        self.blue_team = Team("Blue")
        self.red_team = Team("Red")
        self.n = 0
        if render_mode == "human":
            self.prepare_render()

    @property
    def tanks(self) -> list[Tank]:
        """所有的坦克实例。"""
        return self.blue_team.tanks + self.red_team.tanks

    @property
    def soldiers(self) -> list[Soldier]:
        """所有的士兵实例。"""
        return self.blue_team.soldiers + self.red_team.soldiers

    def assign_actions(self, actions: dict[str, Action]) -> None:
        """按照uid匹配并执行动作。

        :param actions: 一个字典，用来表征坦克的动作，键是坦克的uid，值是动作，由动作类型和动作值两部分组成
        """
        for uid, a in actions.items():
            flag = False
            for obj in self.tanks + self.soldiers:
                if uid == obj.uid:
                    match a[0]:
                        case 0:
                            obj.move(*a[1])
                            flag = True
                        case 1:
                            if isinstance(obj, Tank) is True:
                                obj.launch(*a[1])
                            else:
                                obj.shoot(*a[1])
                            flag = True
                        case _:
                            raise RuntimeError("非法动作类型！")
                if flag is True:
                    break

    def scan_missile_damage(self) -> float:
        r"""扫描当前环境，使坦克导弹伤害生效，并提取奖励。\
        伤害计算公式为：

        .. math:: d = [1-(\frac{r}{R})^2]*damage

        其中 :math:`d` 是对某个坦克造成的伤害， :math:`r` 是该坦克离爆炸中心的距离，:math:`R` 是最远伤害范围，超出该范围收到的伤害则为0。

        :return: 蓝方奖励1
        """
        r1 = 0.0
        for t in self.tanks:
            for m in t.consumed_missiles:
                for e in self.tanks + self.soldiers:
                    if e.alive is False:
                        continue
                    d = compute_distance(m, e)
                    if d < m.radius:
                        damage = (1 - d / m.radius) * m.damage
                        damage = min(damage, e.hp)
                        e.hp -= damage
                        if e.team == "Red":
                            r1 += damage
                        else:
                            r1 -= damage
                        if e.alive is False and isinstance(e, Tank):
                            if e.team == "Red":
                                r1 += 3
                            else:
                                r1 -= 3
        return r1

    def scan_bullet_damage(self) -> float:
        r"""扫描当前环境，使子弹伤害生效，并提取奖励。\
        士兵与子弹轨迹的垂直距离在一定范围d内时，可以狙杀该士兵，否则无法造成伤害：

        .. math:: d = \frac{|Ax_0+By_0+C|}{\sqrt{A^2+B^2}}

        :return: 蓝方奖励2
        """
        r1 = 0.0
        for s in self.soldiers:
            for b in s.consumed_bullets:
                for ss in self.soldiers:
                    if b.owner.uid == ss.uid:
                        continue
                    if ss.alive is False:
                        continue
                    if compute_distance(s, ss) > s.shoot_distance:
                        continue
                    theta1 = atan(s, b)
                    theta2 = atan(s, ss)
                    theta = theta1 - theta2
                    if theta < 0:
                        theta += 2 * math.pi
                    if math.pi < theta < 2 * math.pi:
                        continue
                    d = compute_distance(s, b, ss)
                    if d < 2:
                        ss.hp = 0
                        r1 += ss.hp
        return r1

    def scan_damage(self) -> Reward:
        r"""扫描当前环境，使伤害生效，并提取奖励，包含：

            * 士兵子弹伤害
            * 坦克导弹伤害

        :return: 一个二元组，分别代表蓝队和红队获得的奖励。
        """
        r1 = 0.0
        r1 += self.scan_missile_damage()
        r1 += self.scan_bullet_damage()
        for t in self.tanks:
            t.consumed_missiles.clear()
        for s in self.soldiers:
            s.consumed_bullets.clear()
        if self.red_team.alive is False:
            r1 += 30
        if self.blue_team.alive is False:
            r1 -= 30
        return r1, -r1

    def get_obs(self) -> Observation:
        """获取环境观测信息，包含每个坦克，无论是否死亡。

        :return: 返回每个坦克的（坐标，生命值，速度，导弹剩余量，燃油剩余量）
        """
        obs = {}
        for t in self.tanks:
            obs.update(
                {
                    t.uid: {
                        "x": t.x,
                        "y": t.y,
                        "hp": t.hp,
                        "speed": t.speed,
                        "missile": len(t.missiles),
                        "fuel": t.fuel,
                    }
                }
            )
        for s in self.soldiers:
            obs.update(
                {
                    s.uid: {
                        "x": s.x,
                        "y": s.y,
                        "hp": s.hp,
                        "speed": s.speed,
                        "bullet": len(s.bullets),
                    }
                }
            )
        return obs

    @property
    def terminated(self) -> bool:
        """判断对战是否结束，当一方坦克全部阵亡（即生命值全部不大于0）时，对战结束，返回True并结束当前回合。"""
        return self.blue_team.alive is False or self.red_team.alive is False

    @property
    def truncated(self) -> bool:
        """判断对战是否超出最大回合长度，超出时返回True并结束当前回合。"""
        return self.n >= config.max_steps

    def step(
        self, actions: dict[str, Action]
    ) -> tuple[Observation, Reward, bool, bool, Info]:
        """环境迭代一个时间步，并返回五元组（观测，奖励，终止，截断，额外信息）。

        :param actions: 一个字典，用来表征坦克的动作，键是坦克的uid，值是动作，由动作类型和动作值两部分组成，支持仅传入部分坦克的动作
        """
        self.assign_actions(actions)
        if self.render_mode == "human":
            self.render()
        r = self.scan_damage()
        obs = self.get_obs()
        info = {}
        self.n += 1
        return obs, r, self.terminated, self.truncated, info

    def reset(self) -> tuple[Observation, Info]:
        """重置环境，各个坦克返回原位置，且生命值完整，弹药量完整。"""
        for k in Tank.id_counter:
            Tank.id_counter[k] = 0
        for k in Soldier.id_counter:
            Soldier.id_counter[k] = 0
        self.red_team = Team("Red")
        self.blue_team = Team("Blue")
        self.n = 0
        info = {}
        if self.render_mode == "human":
            self.render()
        return self.get_obs(), info

    def prepare_render(self) -> None:
        """
        准备渲染
            * 初始化画布
            * 设置标题
            * 绘制图例
        """
        _fig, self.window = plt.subplots(figsize=(8, 5))
        _fig.canvas.manager.window.title(
            f"{self.blue_team.n_tank}v{self.red_team.n_tank}坦克对战"
        )
        _fig.subplots_adjust(left=0)
        self.window.set_aspect("equal")
        self.window.axis((0, config.map.width, 0, config.map.height))
        self.window.grid()
        _kwargs = {"xdata": [], "ydata": [], "marker": "o", "ls": "none"}
        __kwargs = {"xdata": [], "ydata": [], "ls": "--"}
        kwargs1 = {**_kwargs, "markeredgecolor": "none"}
        kwargs2 = {**_kwargs, "color": "none"}
        self.window.legend(
            handles=(
                plt.Line2D(
                    **kwargs1, markerfacecolor="r", markersize=12, label="红方坦克"
                ),
                plt.Line2D(
                    **kwargs1, markerfacecolor="b", markersize=12, label="蓝方坦克"
                ),
                plt.Line2D(
                    **kwargs1, markerfacecolor="r", markersize=4, label="红方导弹"
                ),
                plt.Line2D(
                    **kwargs1, markerfacecolor="b", markersize=4, label="蓝方导弹"
                ),
                plt.Line2D(
                    **kwargs2, markersize=6, markeredgecolor="r", label="红方士兵"
                ),
                plt.Line2D(
                    **kwargs2, markersize=6, markeredgecolor="b", label="蓝方士兵"
                ),
                plt.Line2D(**__kwargs, color="r", label="红方子弹"),
                plt.Line2D(**__kwargs, color="b", label="蓝方子弹"),
            ),
            bbox_to_anchor=(1.30, 1.02),
        )

    def render(self) -> None:
        """渲染环境。"""
        scale = self.config.map.width / 50
        for t in self.tanks:
            _alpha = t.hp / config.tank.hp
            _c = t.team.lower()
            if t.alive is True:
                if _alpha < 0.10:
                    _t = Circle(
                        (t.x, t.y), 1 * scale, facecolor="none", edgecolor=_c, ls="--"
                    )
                else:
                    _t = Circle((t.x, t.y), 1 * scale, color=_c, alpha=_alpha)
                _c_f = "w"
            else:
                _t = Circle(
                    (t.x, t.y), 1 * scale, facecolor="none", edgecolor=_c, ls="--"
                )
                _c_f = _c
            self.window.add_patch(_t)
            self.window.text(
                t.x, t.y, t.uid.split(":")[-1], va="center", ha="center", color=_c_f
            )
            if len(self.tanks) <= 10:
                self.window.table(
                    cellText=[
                        (f"{t.hp:.2f}", len(t.missiles), int(t.fuel))
                        for t in self.tanks
                    ],
                    rowLabels=[t.uid for t in self.tanks],
                    colLabels=["生命值", "弹药量", "燃油量"],
                    colWidths=[0.10, 0.10, 0.10],
                    rowLoc="center",
                    cellLoc="center",
                    bbox=[1.30, 0, 0.30, 0.60],
                )
            for m in t.consumed_missiles:
                _m = Circle(
                    (m.x, m.y),
                    config.missile.damage_radius,
                    ls="--",
                    facecolor="None",
                    edgecolor=_c,
                )
                _dot = Circle((m.x, m.y), 0.30 * scale, facecolor=_c)
                self.window.add_patch(_m)
                self.window.add_patch(_dot)
        for s in self.soldiers:
            _alpha = s.hp / config.soldier.hp
            _c = s.team.lower()
            if _alpha < 0.10:
                _s = Circle(
                    (s.x, s.y),
                    0.5 * scale,
                    facecolor="none",
                    edgecolor=_c,
                    ls="--",
                )
            else:
                _s = Circle(
                    (s.x, s.y),
                    0.5 * scale,
                    facecolor="none",
                    edgecolor=_c,
                    alpha=_alpha,
                )
            self.window.add_patch(_s)
            for b in s.consumed_bullets:
                self.window.plot([b.x, s.x], [b.y, s.y], ls="--", color=_c)

        plt.pause(1 / self.render_fps)
        # plt.pause(1)
        for _p in self.window.patches:
            _p.remove()
        for _t in self.window.texts:
            _t.remove()
        for _t in self.window.tables:
            _t.remove()
        for _l in self.window.lines:
            _l.remove()
